# visualizer.py（优化布局 - 控件移至右侧）

import pygame
import sys

# 颜色
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (0, 100, 255)
RED = (255, 0, 0)
GREEN = (0, 200, 0)
YELLOW = (255, 255, 0)
GRAY = (200, 200, 200)
DARK_GRAY = (100, 100, 100)
LIGHT_BLUE = (180, 200, 255)
OBSTACLE = (50, 50, 50)
PATH_COLOR = (0, 255, 100)

from algorithms import dijkstra, a_star


class Visualizer:
    def __init__(self, graph, pos, width=1200, height=800, start_node=0, end_node=99):
        self.graph = graph
        self.pos = pos
        self.width = width
        self.height = height
        self.node_radius = 20
        self.font = None
        self.small_font = None
        self.screen = None
        self.clock = None
        self.start_node = start_node
        self.end_node = end_node
        self.obstacles = set()

        # UI 改为右侧布局
        self.selected_algorithm = "选择算法"
        self.algorithm_options = ["Dijkstra", "A*"]
        self.dropdown_open = False

        # 右侧UI区域
        self.ui_x = 1000  # 控件起始X坐标
        self.ui_width = 180

        self.dropdown_rect = pygame.Rect(self.ui_x, 80, self.ui_width, 40)
        self.option_rects = [pygame.Rect(self.ui_x, 130 + i * 40, self.ui_width, 40) for i in range(len(self.algorithm_options))]
        self.run_button = pygame.Rect(self.ui_x + 30, 350, 120, 50)

        # 状态
        self.dragging = None
        self.placing_obstacle = False

        # ... 原有初始化 ...
        self.shortest_path = []  # 存储最终路径，保持高亮
        self.path_length = 0  # 新增：记录最优路径步数（边数）

    def init_pygame(self):
        pygame.init()
        self.screen = pygame.display.set_mode((self.width, self.height))
        pygame.display.set_caption("100节点网格路径算法可视化")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.SysFont('SimHei', 20)
        self.small_font = pygame.font.SysFont('SimHei', 16)

    def node_at_pos(self, mouse_pos):
        for node, pos in self.pos.items():
            dx = mouse_pos[0] - pos[0]
            dy = mouse_pos[1] - pos[1]
            if dx * dx + dy * dy <= self.node_radius * self.node_radius:
                return node
        return None

    def draw_grid_graph(self, explored=None, updated=None, path=None, current=None):
        for node in range(self.graph.V):
            color = WHITE
            radius = self.node_radius

            if node in self.obstacles:
                color = OBSTACLE
            elif node == self.start_node:
                color = GREEN
                radius = 22
            elif node == self.end_node:
                color = RED
                radius = 22
            elif explored and node in explored:
                color = BLUE
            elif updated and node in updated:
                color = YELLOW
            elif node in self.shortest_path:  # 优先级最高：保持路径高亮
                color = PATH_COLOR
            elif path and node in path:  # 暂时显示的路径（动画中）
                color = PATH_COLOR

            pygame.draw.circle(self.screen, color, self.pos[node], radius)
            pygame.draw.circle(self.screen, BLACK, self.pos[node], radius, 2)

    def draw_ui(self):
        # 绘制右侧控制面板背景
        panel_rect = pygame.Rect(self.ui_x - 10, 50, self.ui_width + 20, self.height - 100)
        pygame.draw.rect(self.screen, (240, 240, 240), panel_rect)
        pygame.draw.rect(self.screen, DARK_GRAY, panel_rect, 2)

        # 标题
        title = self.font.render("路径算法", True, BLACK)
        self.screen.blit(title, (self.ui_x, 20))

        # 下拉框
        bg_color = LIGHT_BLUE
        pygame.draw.rect(self.screen, bg_color, self.dropdown_rect)
        pygame.draw.rect(self.screen, DARK_GRAY, self.dropdown_rect, 2)
        text = self.font.render(self.selected_algorithm, True, BLACK)
        self.screen.blit(text, (self.dropdown_rect.x + 10, self.dropdown_rect.y + 8))

        # 选项
        if self.dropdown_open:
            for i, rect in enumerate(self.option_rects):
                color = WHITE if i % 2 == 0 else LIGHT_BLUE
                pygame.draw.rect(self.screen, color, rect)
                pygame.draw.rect(self.screen, DARK_GRAY, rect, 1)
                text = self.small_font.render(self.algorithm_options[i], True, BLACK)
                self.screen.blit(text, (rect.x + 10, rect.y + 8))

        # 运行按钮
        mouse_pos = pygame.mouse.get_pos()
        btn_color = GREEN if self.run_button.collidepoint(mouse_pos) else (0, 180, 0)
        pygame.draw.rect(self.screen, btn_color, self.run_button)
        pygame.draw.rect(self.screen, DARK_GRAY, self.run_button, 2)
        text = self.font.render("运行", True, WHITE)
        self.screen.blit(text, (self.run_button.x + 25, self.run_button.y + 12))

        # === 使用真实图形绘制说明文字（从 Y=420 开始）===
        y_start = 420
        line_height = 30
        icon_radius = 6
        text_x_offset = 30  # 文字距离图标的偏移

        def draw_legend_item(color, label, y):
            center_y = y + 10
            # 绘制小圆圈
            pygame.draw.circle(self.screen, color, (self.ui_x + 15, center_y), icon_radius)
            pygame.draw.circle(self.screen, BLACK, (self.ui_x + 15, center_y), icon_radius, 1)
            # 绘制文字
            text_surf = self.small_font.render(label, True, DARK_GRAY)
            self.screen.blit(text_surf, (self.ui_x + text_x_offset, y))

        # 操作说明标题
        header1 = self.font.render("操作说明:", True, BLACK)
        self.screen.blit(header1, (self.ui_x, y_start))

        draw_legend_item(GREEN, "设置起点 (左键)", y_start + 25)
        draw_legend_item(RED, "设置终点 (右键)", y_start + 55)
        draw_legend_item(OBSTACLE, "添加障碍 (Shift+左键)", y_start + 85)

        # 算法状态标题
        header2 = self.font.render("算法状态:", True, BLACK)
        self.screen.blit(header2, (self.ui_x, y_start + 135))

        draw_legend_item(BLUE, "探索节点", y_start + 160)
        draw_legend_item(YELLOW, "更新距离", y_start + 190)
        draw_legend_item(PATH_COLOR, "最短路径", y_start + 220)

        # 显示最优路径步数
        if self.path_length > 0:
            steps_text = self.font.render(f"最优路径: {self.path_length} 步", True, BLACK)
        else:
            steps_text = self.font.render("最优路径: -", True, DARK_GRAY)
        self.screen.blit(steps_text, (self.ui_x, y_start + 260))

    def get_algorithm_generator(self, name):
        if name == "Dijkstra":
            return dijkstra(self.graph, self.start_node, self.end_node)
        elif name == "A*":
            return a_star(self.graph, self.start_node, self.end_node, list(self.pos.values()))
        return None

    def run_selection_ui(self):
        self.init_pygame()
        running = True
        animating = False
        gen = None

        final_path_result = []  # 用于保存算法返回的最终路径
        while running:
            self.screen.fill(WHITE)
            self.draw_grid_graph()
            self.draw_ui()
            pygame.display.flip()

            if animating:
                try:
                    action = next(gen)
                    typ, data1, data2 = action
                    self.screen.fill(WHITE)
                    self.draw_grid_graph()
                    self.draw_ui()

                    if typ == 'explore':
                        pygame.draw.circle(self.screen, BLUE, self.pos[data1], self.node_radius)
                    elif typ == 'update':
                        pygame.draw.circle(self.screen, YELLOW, self.pos[data1], self.node_radius)
                    elif typ == 'path' and data1:
                        # data1 是路径节点列表
                        for node in data1:
                            pygame.draw.circle(self.screen, PATH_COLOR, self.pos[node], self.node_radius)
                        # 记录最终路径
                        final_path_result = data1.copy()  # 关键：保存路径
                        # 计算路径边数（节点数 - 1）
                        self.path_length = len(data1) - 1 if data1 else 0

                    pygame.display.flip()
                    pygame.time.wait(100)
                except StopIteration:
                    animating = False
                    gen = None
                    # 保存最终路径
                    self.shortest_path = final_path_result.copy()
                    final_path_result.clear()  # 清空，为下次运行准备

                    # 重绘以显示持久路径
                    self.screen.fill(WHITE)
                    self.draw_grid_graph()
                    self.draw_ui()
                    pygame.display.flip()
                continue

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()

                if event.type == pygame.MOUSEBUTTONDOWN:
                    pos = event.pos
                    node = self.node_at_pos(pos)

                    if node is not None:
                        if event.button == 1:
                            if pygame.key.get_mods() & pygame.KMOD_SHIFT:
                                if node not in (self.start_node, self.end_node):
                                    self.obstacles.add(node)
                                    for i in range(self.graph.V):
                                        self.graph.remove_edge(node, i)
                            else:
                                self.start_node = node
                                if node in self.obstacles:
                                    self.obstacles.remove(node)
                        elif event.button == 3:
                            self.end_node = node
                            if node in self.obstacles:
                                self.obstacles.remove(node)

                    if self.dropdown_rect.collidepoint(pos):
                        self.dropdown_open = not self.dropdown_open
                    else:
                        self.dropdown_open = False
                        for i, rect in enumerate(self.option_rects):
                            if rect.collidepoint(pos):
                                self.selected_algorithm = self.algorithm_options[i]
                                break

                    if self.run_button.collidepoint(pos):
                        if self.selected_algorithm != "选择算法":
                            # 清除上一次的最短路径显示
                            self.shortest_path = []
                            # 强制重绘（清除路径）
                            self.screen.fill(WHITE)
                            self.draw_grid_graph()
                            self.draw_ui()
                            pygame.display.flip()
                            # 启动新算法
                            gen = self.get_algorithm_generator(self.selected_algorithm)
                            if gen:
                                animating = True

        pygame.quit()
